import os, sys
import pandas
import torch

if sys.platform == "darwin":
    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
now_dir = os.getcwd()
sys.path.append(now_dir)

import argparse
from typing import Optional, List
import numpy as np
import ChatTTS
from tools.audio import pcm_arr_to_mp3_view
from tools.logger import get_logger


logger = get_logger("Command")

#########################################################
#文章转为(剧本)
#剧本转音频文件
#合并音频文件（加背景？）
#
#########################################################

def save_mp3_file(wav, index, dir):
    data = pcm_arr_to_mp3_view(wav)
    mp3_filename = f"output_audio_{index}.mp3"
    with open(os.path.join(dir, mp3_filename), "wb") as f:
        f.write(data)
    logger.info(f"Audio saved to {mp3_filename}")


def speaker_seed(num: int):
    """
    设置 CPU 生成随机数的 种子 ，方便下次复现实验结果。
    为 CPU 设置 种子 用于生成随机数，以使得结果是确定的。
    当你设置一个随机种子时，接下来的随机算法生成数根据当前的随机种子按照一定规律生成。
    随机种子作用域是在设置时到下一次设置时。要想重复实验结果，设置同样随机种子即可。
    :param num:
    :return:
    """
    torch.manual_seed(num)

def excel_read():
    """
    将音色配置转换为dict

    :return:
    """
    wb = pandas.read_excel("seed.xlsx",header=0) # header=0 ,第一行作为标题
    data_dict= wb.to_dict()
    return data_dict

def refine_sentence():
    chat = ChatTTS.Chat(get_logger("ChatTTS"))
    chat.RefineTextParams

def main(texts: List[str], spk: Optional[int] = None,stream=False, dir=now_dir):
    """

    :param texts:
    :param spk:  seed #2222、7869、6653是男性音色，3333、4099、5099 是女性角色
    :param stream:
    :return:
    """
    logger.info("Text input: %s", str(texts))
    print("spk is",spk)
    chat = ChatTTS.Chat(get_logger("ChatTTS"))
    logger.info("Initializing ChatTTS...")
    if chat.load(source="custom", compile=False, custom_path="./asset"):
        logger.info("Models loaded successfully.")
    else:
        logger.error("Models load failed.")
        sys.exit(1)

    if spk == None:
        spk_emb = chat.sample_random_speaker()
    else:
        if isinstance(spk,int):
            speaker_seed(spk)
            spk_emb = chat.sample_random_speaker()
        else:
            # sample_random_speaker 返回的是一个 768 维的向量，也就是说如果你把这个 768 维的向量固定住，就可以获取一个固定的音色
            spk_emb = spk
    logger.info("Use speaker:")

    params_infer_code = ChatTTS.Chat.InferCodeParams(
        spk_emb=spk_emb,  # add sampled speaker
        temperature=.6,  # using custom temperature
        top_P=1,  # top P decode
        top_K=30,  # top K decode
        prompt='[speed_0]')
    # For sentence level manual control.
    # use oral_(0-9), laugh_(0-9), break_(0-7)
    # to generate special token in text to synthesize.
    params_refine_text = ChatTTS.Chat.RefineTextParams(
        prompt='[oral_0][laugh_0][break_6]',
    )

    logger.info("Start inference.")
    wavs = chat.infer(
        texts,
        stream,
        skip_refine_text=True,
        refine_text_only=False,
        params_refine_text=params_refine_text,
        params_infer_code=params_infer_code
    )
    logger.info("Inference completed.")
    # Save each generated wav file to a local file
    if stream:
        wavs_list = []
    for index, wav in enumerate(wavs):
        if stream:
            for i, w in enumerate(wav):
                save_mp3_file(w, (i + 1) * 1000 + index, dir)
            wavs_list.append(wav)
        else:
            save_mp3_file(wav, index, dir)
    if stream:
        for index, wav in enumerate(np.concatenate(wavs_list, axis=1)):
            save_mp3_file(wav, index, dir)
    logger.info("Audio generation successful.")


if __name__ == "__main__":
    logger.info("Starting ChatTTS commandline demo...")
    # parser = argparse.ArgumentParser(
    #     description="ChatTTS Command",
    #     usage='[--spk xxx] [--stream] "Your text 1." " Your text 2."',
    # )
    # parser.add_argument(
    #     "--spk",
    #     help="Speaker (empty to sample a random one)",
    #     type=Optional[str],
    #     default=None,
    # )
    # parser.add_argument(
    #     "--stream",
    #     help="Use stream mode",
    #     action="store_true",
    # )
    # parser.add_argument(
    #     "texts",
    #     help="Original text",
    #     default=["YOUR TEXT HERE"],
    #     nargs=argparse.REMAINDER,
    # )
    # args = parser.parse_args()
    # main(args.texts, args.spk, args.stream)
    # logger.info("ChatTTS process finished.")
    texts = [
        '“呵呵哈轰。这是一个测试。你知道吗？']
    main(texts, 2, False)
    # print(len("呵呵"))
